from core.function import Function
from core.cuda import cuda


class Pow(Function):
    def __init__(self, c):
        """
        初始化 Pow 操作。

        Parameters:
            c (float): 指数。
        """
        super().__init__()
        self.c = c

    def forward(self, x):
        """
        执行 Pow 操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: Pow 操作的结果。
        """
        x = cuda.to_array(x)
        y = x ** self.c
        return y

    def backward(self, gy):
        """
        执行 Pow 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """
        x, = self.inputs
        c = self.c

        gx = c * x ** (c - 1) * gy
        return gx

    def pow(self, x):
        """
        执行 Pow 操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: Pow 操作的结果。
        """
        return self(x)

